package 剑指offer;

import org.junit.Assert;

import java.util.PriorityQueue;

/**
 * @descripition:
 * @author: yangwk37982
 * @date: 2021/7/3 9:09
 */
public class 查找中位数 {
    public static class MidQueue<T extends Comparable> {
        PriorityQueue<T> maxQueue; //顶端存比中位数小的数
        PriorityQueue<T> minQueue;//顶端存中位数大的数

        public MidQueue() {
            maxQueue = new PriorityQueue<T>((o1, o2) -> -o1.compareTo(o2));
            minQueue = new PriorityQueue<T>((o1, o2) -> o1.compareTo(o2));
        }

        public void insert(T ele) {
            if (maxQueue.isEmpty()) {
                maxQueue.add(ele);
            } else if (minQueue.isEmpty()) {
                if (ele.compareTo(maxQueue.peek()) >= 0) {
                    minQueue.add(ele);
                } else {
                    minQueue.add(maxQueue.poll());
                    maxQueue.add(ele);
                }
            } else {
                T max = maxQueue.peek();
                T min = minQueue.peek();
                if (maxQueue.size() == minQueue.size()) {
                    if (ele.compareTo(min) >= 0) {
                        minQueue.add(ele);
                    } else {
                        maxQueue.add(ele);
                    }

                } else if (maxQueue.size() > minQueue.size()) {
                    if (ele.compareTo(max) >= 0) {
                        minQueue.add(ele);
                    } else {
                        minQueue.add(maxQueue.poll());
                        maxQueue.add(ele);
                    }
                } else {
                    if (ele.compareTo(min) <= 0) {
                        maxQueue.add(ele);
                    } else {
                        maxQueue.add(minQueue.poll());
                        minQueue.add(ele);
                    }
                }
            }
        }

        public double getMid() {
            if (maxQueue.size() == minQueue.size()) {
                return ((Integer) maxQueue.peek() + (Integer) minQueue.peek()) / 2.0;
            } else if (maxQueue.size() > minQueue.size()) {
                return  maxQueue.peek() == null ?  (Integer)minQueue.peek():(Integer)maxQueue.peek() ;
            }
            return minQueue.peek() == null ? (Integer) maxQueue.peek() : (Integer) minQueue.peek();
        }
    }
    public static void main(String[] args) {
        MidQueue midQueue = new MidQueue();
        Integer[] arr = new Integer[]   {5, 2, 1, 6, 5, 0, 0, 1, 0};
        double[] arrEqual = new double[]{5.0, 3.5, 2.0, 3.5, 5.0, 3.5,2.0,1.5,1.0};
        for (Integer i = 0; i < arr.length; i++) {
            midQueue.insert(arr[i]);
            Assert.assertEquals(String.valueOf(midQueue.getMid()), String.valueOf(arrEqual[i]));
        }

    }
}
